package com.example.netcapture.utils;

import com.example.netcapture.entity.PacketInfo;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class TcpSessionManager {
    private static final Map<String, TcpSession> sessions = new ConcurrentHashMap<>();

    public static TcpSession getOrCreateSession(String sessionId) {
        return sessions.computeIfAbsent(sessionId, k -> new TcpSession());
    }

    public static class TcpSession {
        private final List<PacketInfo> packets = new ArrayList<>();
        private final List<FileDataChunk> fileDataChunks = new ArrayList<>();
        private long expectedSeqClient = 0;
        private long expectedSeqServer = 0;
        private String currentFileType = null;
        private long fileStartSeq = -1;

        public void addPacket(PacketInfo packet) {
            // 按序列号插入到正确位置，避免重复
            synchronized (packets) {
                boolean inserted = false;
                for (int i = 0; i < packets.size(); i++) {
                    if (packets.get(i).getTcpSeqNumber() > packet.getTcpSeqNumber()) {
                        packets.add(i, packet);
                        inserted = true;
                        break;
                    }
                }
                if (!inserted) {
                    packets.add(packet);
                }
            }
        }

        /**
         * 添加文件数据块
         */
        public void addFileData(byte[] data, String fileType) {
            synchronized (fileDataChunks) {
                FileDataChunk chunk = new FileDataChunk(data, fileType, System.currentTimeMillis());
                fileDataChunks.add(chunk);

                // 如果文件类型改变，重置文件重组状态
                if (currentFileType == null || !currentFileType.equals(fileType)) {
                    currentFileType = fileType;
                    fileStartSeq = -1;
                }

                logFileDataChunk(chunk);
            }
        }

        /**
         * 添加文件数据块（带序列号信息）
         */
        public void addFileData(byte[] data, String fileType, long sequenceNumber) {
            synchronized (fileDataChunks) {
                FileDataChunk chunk = new FileDataChunk(data, fileType, sequenceNumber, System.currentTimeMillis());
                fileDataChunks.add(chunk);

                // 记录文件开始的序列号
                if (fileStartSeq == -1) {
                    fileStartSeq = sequenceNumber;
                }

                // 如果文件类型改变，重置文件重组状态
                if (currentFileType == null || !currentFileType.equals(fileType)) {
                    currentFileType = fileType;
                    fileStartSeq = sequenceNumber;
                }

                logFileDataChunk(chunk);
            }
        }

        /**
         * 获取重组后的完整文件数据
         */
        public byte[] getReassembledFileData() {
            synchronized (fileDataChunks) {
                if (fileDataChunks.isEmpty()) {
                    return new byte[0];
                }

                // 计算总大小
                int totalSize = fileDataChunks.stream().mapToInt(chunk -> chunk.data.length).sum();
                byte[] result = new byte[totalSize];

                // 按时间戳排序并合并数据
                fileDataChunks.stream()
                        .sorted((c1, c2) -> Long.compare(c1.timestamp, c2.timestamp))
                        .forEach(chunk -> {
                            int currentLength = getCurrentLength(result);
                            System.arraycopy(chunk.data, 0, result, currentLength, chunk.data.length);
                        });

                return result;
            }
        }

        /**
         * 获取当前有效数据长度
         */
        private int getCurrentLength(byte[] data) {
            for (int i = 0; i < data.length; i++) {
                if (data[i] == 0) {
                    return i;
                }
            }
            return data.length;
        }

        /**
         * 获取文件信息统计
         */
        public FileSessionInfo getFileSessionInfo() {
            synchronized (fileDataChunks) {
                FileSessionInfo info = new FileSessionInfo();
                info.fileType = currentFileType;
                info.chunkCount = fileDataChunks.size();
                info.totalSize = fileDataChunks.stream().mapToInt(chunk -> chunk.data.length).sum();
                info.startTime = fileDataChunks.stream()
                        .mapToLong(chunk -> chunk.timestamp)
                        .min()
                        .orElse(System.currentTimeMillis());
                info.endTime = fileDataChunks.stream()
                        .mapToLong(chunk -> chunk.timestamp)
                        .max()
                        .orElse(System.currentTimeMillis());
                return info;
            }
        }

        /**
         * 清除文件数据
         */
        public void clearFileData() {
            synchronized (fileDataChunks) {
                fileDataChunks.clear();
                currentFileType = null;
                fileStartSeq = -1;
            }
        }

        /**
         * 记录文件数据块日志
         */
        private void logFileDataChunk(FileDataChunk chunk) {
            System.out.printf("File data chunk added - Type: %s, Size: %d bytes, Seq: %d%n",
                    chunk.fileType, chunk.data.length, chunk.sequenceNumber);
        }

        public String getReassembledStream() {
            // 按序列号排序并重组数据
            return TcpStreamReassembler.reassembleTcpData(packets);
        }
    }

    /**
     * 文件数据块类
     */
    public static class FileDataChunk {
        public byte[] data;
        public String fileType;
        public long sequenceNumber;
        public long timestamp;

        public FileDataChunk(byte[] data, String fileType, long timestamp) {
            this.data = data;
            this.fileType = fileType;
            this.sequenceNumber = -1;
            this.timestamp = timestamp;
        }

        public FileDataChunk(byte[] data, String fileType, long sequenceNumber, long timestamp) {
            this.data = data;
            this.fileType = fileType;
            this.sequenceNumber = sequenceNumber;
            this.timestamp = timestamp;
        }
    }

    /**
     * 文件会话信息类
     */
    public static class FileSessionInfo {
        public String fileType;
        public int chunkCount;
        public int totalSize;
        public long startTime;
        public long endTime;

        @Override
        public String toString() {
            return String.format("FileSessionInfo{fileType='%s', chunkCount=%d, totalSize=%d, duration=%dms}",
                    fileType, chunkCount, totalSize, (endTime - startTime));
        }
    }

    /**
     * 获取所有会话的文件传输统计
     */
    public static Map<String, FileSessionInfo> getFileTransferStatistics() {
        Map<String, FileSessionInfo> statistics = new ConcurrentHashMap<>();
        sessions.forEach((sessionId, session) -> {
            FileSessionInfo info = session.getFileSessionInfo();
            if (info.chunkCount > 0) {
                statistics.put(sessionId, info);
            }
        });
        return statistics;
    }

    /**
     * 根据会话ID获取文件数据
     */
    public static byte[] getFileDataBySession(String sessionId) {
        TcpSession session = sessions.get(sessionId);
        if (session != null) {
            byte[] fileData = session.getReassembledFileData();
            if (fileData != null && fileData.length > 0) {
                return fileData;
            }
        }
        return null;
    }

    /**
     * 清理空闲会话（可选功能）
     */
    public static void cleanupIdleSessions(long timeoutMs) {
        long currentTime = System.currentTimeMillis();
        sessions.entrySet().removeIf(entry -> {
            TcpSession session = entry.getValue();
            FileSessionInfo info = session.getFileSessionInfo();
            return (currentTime - info.endTime) > timeoutMs && info.chunkCount == 0;
        });
    }
}